#!/usr/bin/python 
###########################################################
# PuSSH - Pythonized ubiquitous SSH # Version 2 ###########
###########################################################
# an Open Source script originally made at CERN (2003) ####
###########################################################
# by A. Doicin                      #######################
###########################################################
# This program is free software; you can redistribute it 
# and/or modify it under the terms of the GNU General 
# Public License as published by the Free Software 
# Foundation; either version 2 of the License, or (at your 
# option) any later version.
###########################################################
# http://pussh.sourceforge.net
###########################################################

import os
import sys
import getopt
import select
import string

###########################################################
# Function "Usage" - self-explanatory really ... and also 
# a decent synopsis of script functionality.
###########################################################

def Usage():
    """
 PUSSH is "Pythonic Ubiquitous SSH" - a big SSH wrapper consisting of two Python modules (at this time one is called explicitly by the other, although this may change in future versions) - for usage on networks with many/multiple hosts, ideally wherein SSH is configured with Kerberos, RSA/DSA keys, or SSH-Agent in such a way as to avoid any actual password authentication on the command line. Using PuSSH, you can send the same command via SSH to a range of machines of any size.

 ---

 PUSSH Usage:

 pussh [-h] [-s] [-r] [-f <filename>] [-p <parallel>] [-P <port>] [-l <userid>] [-t <timeout value in seconds>] <target> <commands>

 -h                             print this helpful message

 -p                             parallel - numeric argument is equal to number of machines on which command 
								should execute in parallel. If NOT specified, this option will be SET by 
								default, and the size of the "parallel slice" is set to a finely tuned 42. 

 -P                             a numerical value indicating the TCP port on which the target machine's SSHD 
								server is (or should be) listening. If NOT specified, this value is SET to 22 
								by default.

 -l <user>                      send command as user <user> ; if NOT specified, this option is SET to the user 
								who owns the current shell.

 -r                             send command as root (a lazy way of specifying -l root ... see above).

 -t <timeout value in seconds>  time after which a command times out for any particular reason; if NOT 
								specified, this option is SET to 60 seconds by default.

 -f <filename>                  file contains list of target-hosts, each on one row of the file.

 -s                             execute a command over a range of machines in SEQUENCE (PARALLEL execution is 
                                the default). This can be useful when a password is necessary on one or more 
                                machines in a range or list for some reason.

 --prefix                       prefix each line of output with 'hostname:' 
                                (note: this option is rather crudely implemented and so output may not always 
								be 100% satisfactory - please feel free to voice your complaints).
     
 <target(s)>                    target machine(s) in the following possible formats:

                                machine                 - single target
                                machine[01-10]          - range delimited by numeric(-numeric) suffixes in 
														  square brackets
                                machine[01,05,10]       - range delimited by numeric(,numeric) suffixes 
														  separated by commas
                                machine[01-05,07,09-10] - any combination or permutation of the above 2

 ---

 Example 1:                     pussh -r machine[01-10,12,15-20] uptime 

 Example 2:                     pussh machine[01-10,12,15-20] "uname -a; uptime"

 Example 3:                     pussh -t 10 -f hostlist.txt "uname -a; uptime; ls -l | grep stuff"
 
 Example 4:                     pussh -t 10 -P 55 -f hostlist -r "uname -a; uptime; ls -l | grep stuff"
"""

###########################################################
# More practical Function definitions start here ... 
###########################################################
def InvalidTargetting():
    print "PUSSH(001): Invalid host targetting - do try try again."
    os._exit(1)
###########################################################
def test_for_numeric(var): 
    try: var = int(var)
    except ValueError:
        print "PUSSH(001): Non-numerical variable:", var, "in timeout value (-t xx), range suffix, or parallel group spec (-p xx)."
        os._exit(1)
    return
###########################################################
def DeterminePath(login):
    progpath = os.path.dirname(sys.argv[0])
    return progpath
###########################################################
def optionizer():
    """PUSSH options specification:
 
    pussh [-h] [-s] [-r] [-f <filename>] [-P <sshd TCP port>] [-l <userid>] [-t <timeout value in seconds>] <target> <commands>"""

    unique_case = 0                   # by default, we don't expect a scenario wherein the only option is -s
                                      # note: this entire function could be improved by turning it into a dictionary
                                      # ... too many "special cases" may not be good programming form, but for now, what the heck.
                                        
    if loginspec == 1 and timeoutspec == 0 and parallel == 1: 
        opty = "-P ", port, "-l ", login
        # print "loginspec is 1, timeoutspec is 0, and parallel is 1"
    elif loginspec == 1 and timeoutspec == 0 and parallel == 0: 
        opty = "-P ", port, "-l ", login, " -s "
        # print "loginspec is 1, timeoutspec is 0, and parallel is 0"
    elif loginspec == 1 and timeoutspec == 1 and parallel == 1: 
        opty = "-P ", port, "-l ", login, " -t ", timevar
        # print "loginspec is 1, timeoutspec is 1, and parallel is 1"
    elif loginspec == 1 and timeoutspec == 1 and parallel == 0: 
        opty = "-P ", port, "-l ", login, " -t ", timevar, " -s "
        # print "loginspec is 1, timeoutspec is 1, and parallel is 0"
    elif loginspec == 0 and timeoutspec == 1 and parallel == 1: 
        opty = "-P ", port, "-t ", timevar
        # print "loginspec is 0, timeoutspec is 1, and parallel is 1"
    elif loginspec == 0 and timeoutspec == 1 and parallel == 0: 
        opty = "-P ", port, "-t ", timevar, " -s "
        # print "loginspec is 0, timeoutspec is 1, and parallel is 0"
    elif loginspec == 0 and timeoutspec == 0 and parallel == 0:
        # print "loginspec is 0, timeoutspec is 0, and parallel is 0"
        opty = opto = "-s"
        unique_case = 1
    else: 
        # print "apparently nothing, but our port is ...", port
        opty = "-P ", port
    if unique_case == 0: opto = string.join(opty,' ')           # heavily crucial
    return opto, unique_case
############################################################
def range_by_dash(x,y):
    list = range(x,y)
    return list
############################################################
def RangeKruncha(target):
    lsbp = string.find(target, "[")
    if lsbp == 0: InvalidTargetting()
    target_prefix = string.join(target[:lsbp],'')
    # print "target prefix is", target_prefix
    target_suffix =  string.join(target[lsbp+1:target_length-1],'')
    # print "target suffix is", target_suffix
    if len(target_suffix) < 3: InvalidTargetting()
    target_list = []
    dashcount = string.count(target_suffix,"-")
    commacount = string.count(target_suffix,",")
    if dashcount == commacount == 0: InvalidTargetting()
    # check for commas ",", dashes "-" and stuff ...
    if string.count(target_suffix,",")>0:
        suffix_list = string.split(target_suffix,",")
        for arg in suffix_list:
            ranger = 0
            # print "testing", arg
            if string.count(arg,"-") == 1:
                range_limits = string.split(arg,"-")
                for i in range_limits:
                    test_for_numeric(i)
                    if range_limits[0] < range_limits[1] and len(range_limits[0]) == len(range_limits[1]):
                        ranger = 1
                        lenny = len(range_limits[0])
                        for limit in range_limits: test_for_numeric(limit)
                        lowerlimit = int(range_limits[0])
                        upperlimit = int(range_limits[1])
                    else: InvalidTargetting()
            elif string.count(arg,"-")>1:
                print "PUSSH(001): ... range is invalid on target component(s) ...", arg
                InvalidTargetting()
            else:
                test_for_numeric(arg)
                # print arg
                lenny = len(arg)
                next_target_list = [int(arg)]
                target_list = target_list + next_target_list
            if ranger == 1:
                # print "range aii is", range_by_dash(lowerlimit, upperlimit+1)
                next_target_list = range_by_dash(lowerlimit, upperlimit+1)
                target_list = target_list + next_target_list
    elif string.count(target_suffix,"-")==1 and string.count(target_suffix,",")==0:
        range_limits = string.split(target_suffix,"-")
        for limit in range_limits: test_for_numeric(limit)
        if range_limits[0] < range_limits[1] and len(range_limits[0]) == len(range_limits[1]):
            lenny = len(range_limits[0])
            lowerlimit = int(range_limits[0])
            lenny = len(range_limits[0])
            # print "lower limit on singular range is", lowerlimit
            upperlimit = int(range_limits[1])
            # print "upper limit on singular range is", upperlimit
            for i in range_limits:
                # print "testing", i
                test_for_numeric(i)
            # print "singular range is", range_by_dash(lowerlimit, upperlimit+1)
            next_target_list = range_by_dash(lowerlimit, upperlimit+1)
            target_list = target_list + next_target_list
        else:
            print "PUSSH(001): ... invalid range construction ..."
            InvalidTargetting()

    elif string.count(target_suffix,"-")>1 and string.count(target_suffix,",") == 0: InvalidTargetting()

    hostlist = []
    for host in target_list:
        host = target_prefix + string.zfill(host,lenny)
        # print string.zfill(host,lenny)
        hostlist.insert(len(hostlist), host)                     # hostlist is finally ready
    return hostlist                                              # ... return to sender ...
############################################################
def keyboard_interrupt():
    print "PUSSH(001): <<<<< _Clean_Keyboard_Interrupt_  >>>>>"
    os._exit(1)
############################################################
def sequentialCommand(host):
    letsgo = '%s/brussh.py %s %s %s 2>&1' % (progpath, opto, host, payload)
    try:
        os.system('echo %s:' % (host))
        os.system(letsgo)
    except KeyboardInterrupt: print keyboard_interrupt()
############################################################
def parallelCommand(host):
    while prefix == 1: 
        return "%s/brussh.py %s %s %s 2>&1 | sed 's/^/%s: /'" % (progpath, opto, host, payload, host)
    else: return "echo %s: ; %s/brussh.py %s %s %s 2>&1" % (host, progpath, opto, host, payload)
############################################################
def readout():
    while 1:
        try:
            # Could put a timeout here if we had something else to do ...
            readable,writable,errors=select.select(readPipes,[],[])
            for p in readable:
                print p.read().rstrip()
                readPipes.remove(p)
                # os.wait() # Don't want zombies
            if len(readPipes)==0: break
        except KeyboardInterrupt: print keyboard_interrupt()
############################################################
def SliceyNicey(hostlist,step,originalstep,SliceItNice):         # parallel = 1 ... easy life, or so you might think ...
    stepsize = int(originalstep)
    hostlistlen = int(len(hostlist))
    if hostlistlen > stepsize:
        SliceItNice = 1
        slices, remainder = divmod(hostlistlen,stepsize)

        ### Firstly loop through each 'slice of however big the slice is'

        for i in range(slices):
            # print "step is now", hostlist[step:step+stepsize]
            for host in hostlist[step:step+stepsize]: readPipes.append(os.popen(parallelCommand(host)))
            readout()
            step = step + stepsize

    ### and then through the remainder, if there is one.

    if SliceItNice == 1 and remainder == 0: os._exit(0)
    elif SliceItNice == 1 and remainder > 0:
        for host in hostlist[step:]: readPipes.append(os.popen(parallelCommand(host)))
        readout()

    ### If target list is less than whatever the default slice-size is, fire 
    ### off commands to entire target list in parallel.

    elif SliceItNice == 0:
        for host in hostlist: readPipes.append(os.popen(parallelCommand(host)))
        readout()
    return

##################### END of FUNCTIONs #####################

################### Start of MAIN PROGRAM ##################

if __name__ == "__main__":
############################################################
# Declaration of global variables and defaults
############################################################

    timeoutspec = 0      # this means the default timeout will
                         # kick in on "brush" which is 60 secs
                         # i.e. assumes timeout is not
                         # specified on the command line

    loginspec = 0        # this means the default user will
                         # kick in on "brussh" which is the
                         # current user i.e. assumes login is
                         # not specified on the command line
                     
    login = os.environ['LOGNAME']         # as per above

    port = "22"

    progpath = DeterminePath(login)

    # a few defaults to start, concerned with parallelism ######

    parallel = 1         # execute commands in parallel (0 will
                         # mean sequential processing)

    sequence = 0         # means the same as above. but in reverse

    readPipes = []       # empty "piperack" to start

    SliceItNice = 0      # no need to slice, i.e. range or list
                         # of hosts does not exceed 40 or 
                         # specified range via option -p with argument 

    prefix = 0           # do not prefix output with 'hostname: '

    step = 0             # initialize
    originalstep = 42    # default no. of machines to 
                         # execute parallel commands on

####### otherwise, assume nothing ...  #####################

    target_is_singular = target_is_range = target_is_file = 0

############################################################
# Command line options pre-amble ...
############################################################
    try: 
        options, arguments = getopt.getopt(sys.argv[1:],'P:l:t:p:f:rsh',['port','login','timeout','parallel','file','root','sequence','help','prefix','prefixitright'])
    except getopt.error:
        print Usage.__doc__
        os._exit(1)

    for opt in options:
        if opt[0] == '-t':
            timeoutspec = 1
            timevar = opt[1]
            test_for_numeric(timevar)
            if int(timevar) < 301 and int(timevar) > 0: timeout_value = int(timevar)
            else:
                print "Invalid timeout value - timeout must be" 
                print "greater than or equal to 1 second, and" 
                print "less than or equal to 300 seconds."
                print "Type pussh -h <ENTER> for help."
                os._exit(1)

        if opt[0] == '-f':
            target_is_file = 1
            file = opt[1]
 
        if opt[0] == '-s': 
            parallel = 0
            sequence = 1
    
        if opt[0] == '-p':
            stepspec = parallel = 1
            size_of_slice = opt[1]
            test_for_numeric(size_of_slice)
            originalstep = size_of_slice                    

        if opt[0] == '-P':
            portspec = 1
            port = opt[1]

        if opt[0] == '-r':
            loginspec = 1
            login = "root"
            # print "login now=", login

        if opt[0] == '-l':
            loginspec = 1
            login = opt[1]

        if opt[0] == '--prefix': 
            arglist = sys.argv[1:]
            arglist.remove('--prefix')
            arglist = string.join(arglist)
            os.system('%s/pussh --prefixitright %s 2>&1 | sort -k 1 -n' % (progpath,arglist)) 
            os._exit(0)       
 
        if opt[0] == '--prefixitright': 
            prefix =1
        
        if opt[0] == '-h': 
            print Usage.__doc__
            os._exit(1)

    if len(arguments) < 1 and target_is_file:
        print "PUSSH(001) Error: not enough command to send to hostlist in file \"" + file + "\" - you need at least 1 command. Type pussh -h <ENTER> for help."
        os._exit(1)

    if len(arguments) < 2 and target_is_file == 0:
        print "PUSSH(001) Error: not enough command to send to host - you need at least 1 command. Type pussh -h <ENTER> for help."
        os._exit(1)

    else: pass

    if (parallel == sequence) or (parallel == sequence): 
        print "PUSSH(001) Error: Options are incorrect - you cannot specify parallel and sequential processing concurrently."
        print optionizer.__doc__
        os._exit(1)

### and of options preamble, and start of 'what is the nature of the target?' pre-amble ##############################

    try: target = arguments[0]
    except IndexError:
        print "PUSSH(001) Error: no target specified."
        os._exit(1)

    # assume no square bracketed suffix i.e. no multi-targetting 

    target_length = len(target)

    # find no. of left/right square brackets or "lsbs" and "rsbs" ... there can be only 1 of each ...
    # if we find some, then there's a range of targets rather than a single machine target ...
    # which means target_is_range = 1

    lsbs = string.count(target,"[")
    rsbs = string.count(target,"]")

    if lsbs == rsbs == 1 and target[target_length-1] == "]": target_is_range = 1
    
    elif lsbs == rsbs == 0: target_is_singular = 1         # note: this could also mean a single file containing a
                                                           # list of target hosts ... 
    else: InvalidTargetting()

#######################################################################
### where the action really begins - and ends ...                ######
#######################################################################

    if target_is_range: 
    
        hostlist = RangeKruncha(target)

        # print "target list is", target_list
        unique_case = 0
        opto = optionizer()[0]
        payload = '\"' + string.join(arguments[1:]) + '\"'

        if parallel == 0 and target_is_range:
            for host in hostlist: sequentialCommand(host)            # no parallelism, as per option -s for sequential op
            os._exit(0)

        # else ...

        ### If target list exceeds slice size, target list into groups of
        ### "slice-size", so we only ever fire off that many ssh commands
        ### at one time.    

        SliceyNicey(hostlist,step,originalstep,SliceItNice)          # parallelism kicks in

#######################################################################    
    elif target_is_singular == 1 and target_is_file == 0:            # i.e. target is NOT a file ... but a SINGLE host
        host = target
        payload = '\"' + string.join(arguments[1:]) + '\"'

        opto = optionizer()[0]                                       # sort out these troublesome options ...
        if parallel == 0:
            sequentialCommand(host)                                  # no parallelism, as per option -s for sequential op
            os._exit(0)

        # else

        readPipes.append(os.popen(parallelCommand(host)))            # parallelism kicks in
        readout()

#######################################################################
    elif target_is_singular and target_is_file:                      # there can only be a single file, with 1 host per line
        hostlist = []
        f = open(file)
        hostlist = f.read().splitlines()
        payload = '\"' + string.join(arguments[0:]) + '\"'
        opto = optionizer()[0]
    
        if optionizer()[1]:                                          # by implication, parallel = 0 but unique_case also  1 (see optionizer()) ...
            for host in hostlist: sequentialCommand(host)            # ... something to tidy up here ? ... wtf =)
            os._exit(0)
    
        if optionizer()[1] == parallel == 0:                         # parallel = 0, but unique case = 0, so other options are specified also ...
            for host in hostlist: sequentialCommand(host)
            os._exit(0)

        else: SliceyNicey(hostlist,step,originalstep,SliceItNice)    # parallel = 1 ...
